import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np

# 设置随机种子以确保结果可复现
torch.manual_seed(42)
np.random.seed(42)

# 模拟数据生成
def generate_dummy_data(num_samples=1000):
    # 输入是64x64的矩阵，展平为4096维向量
    X = np.random.rand(num_samples, 64, 64).reshape(num_samples, -1)
    # 输出是4类标签，随机生成
    y = np.random.randint(0, 4, size=num_samples)
    return X, y

# 生成训练和测试数据
X_train, y_train = generate_dummy_data(1000)
X_test, y_test = generate_dummy_data(200)

# 转换为PyTorch张量
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.long)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.long)

# 创建数据加载器
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 输入大小是64x64=4096，隐藏层大小为256，输出大小为4
input_size = 64 * 64
hidden_size = 256
output_size = 4

# 初始化模型、损失函数和优化器
model = SimpleNet(input_size, hidden_size, output_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    for batch_X, batch_y in train_loader:
        # 前向传播
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # 每个epoch打印一次损失
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

# 测试模型
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for batch_X, batch_y in test_loader:
        outputs = model(batch_X)
        _, predicted = torch.max(outputs.data, 1)
        total += batch_y.size(0)
        correct += (predicted == batch_y).sum().item()
    
    print(f"Accuracy on test data: {100 * correct / total:.2f}%")

# 测试一个样本
sample_input = X_test_tensor[0].unsqueeze(0)  # 添加批次维度
with torch.no_grad():
    output = model(sample_input)
    _, predicted = torch.max(output.data, 1)
    print(f"Sample input predicted class: {predicted.item()}")